# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import json
import logging

from absl import app, flags

from coin.experimental.xguo.fruit.util.fee import get_fee_rate
from coin.experimental.xguo.fruit.grape.price_graph import PriceGraph
from coin.experimental.xguo.fruit.base.event_clock import event_clock_scheduler, event_clock
from coin.experimental.xguo.fruit.base.strategy_sim_base import (
    BaseSimStrategy,
    BaseStrategyConfig,
)
from coin.strategy.mm.simple_sim.executors import (
    PassAggSimExecutor,
    AggressionSimExecutor,
)

FLAGS = flags.FLAGS


class TriArbSimStrategy(BaseSimStrategy):
  def __init__(self, config):
    super().__init__(config)

    pivot = None
    for currency in self.currency_list:
      if currency.currency == 'USDT':
        pivot = [currency]
        break
    self._pricer = PriceGraph(self.product_list, pivot)

    self.amount_in_usdt = 1000
    self.last_fire_time = 0
    self.pnl_report = []

    self._agg_executor = {}
    for product in self.product_list:
      fee_bp = get_fee_rate(product)
      maker_fee = fee_bp[0] * 1e-4
      taker_fee = fee_bp[1] * 1e-4
      self._agg_executor[product] = AggressionSimExecutor(
          maker_fee=maker_fee,
          taker_fee=taker_fee,
          execution_delay=0,
          max_qty=100,
      )

    self._pnl = 0
    self._vol = 0
    event_clock_scheduler.add_periodic_task(1800, self.log_sim_time)
    event_clock_scheduler.add_periodic_task(300, self.pnl_snapshot)
    event_clock_scheduler.add_periodic_task(3600, self.save_sim_result)

  @property
  def executors(self):
    return self._agg_executor

  def on_book(self, product, book):
    timestamp = book.timestamp
    self._pricer.update(book)
    self._agg_executor[book.product].update_book(book)

    path_info = self._pricer.arb_chance(profit_bp=1)
    if not path_info:
      return

    if timestamp - self.last_fire_time < 5e9:
      return

    for executor in self._agg_executor.values():
      executor.cancel_all(timestamp)

    print(path_info)
    source_qty = self.amount_in_usdt
    for edge_info in path_info.edge_infos:
      order = edge_info.gen_order(source_qty=source_qty)
      source_qty = order.target_qty
      self._agg_executor[order.product].submit(
          book.timestamp,
          order.price,
          order.qty,
          order.order_side,
      )
    self.last_fire_time = timestamp
    self._logger.info('Submit order: %s', event_clock.now())
    self._pnl += path_info.weight
    self._vol += 1
    self._logger.info('pnl: %.4f, %f', self._pnl, self._vol)

  def on_trade(self, product, trade_event):
    self._agg_executor[product].update_trade(trade_event)

  def pnl_snapshot(self):
    timestamp = event_clock.timestamp
    result = {
        'timestamp': timestamp,
        'human_readable_timestamp': str(event_clock.now()),
        'data': {},
    }
    for product, executor in self.executors.items():
      result['data'][product.symbol] = executor.get_summary()
    print(json.dumps(result, indent=2))
    self.pnl_report.append(result)

  def save_sim_result(self):
    filename = self.strategy_config.sim_config.get('result_file')
    if not filename:
      filename = str(event_clock.timestamp) + '.json'
    with open(filename, 'w') as out_file:
      json.dump(self.pnl_report, out_file)


def main(_):
  config = BaseStrategyConfig.from_file(FLAGS.config)
  strategy = TriArbSimStrategy(config)
  strategy.run()
  strategy.save_sim_result()


if __name__ == '__main__':
  flags.DEFINE_string('config', None, 'Strategy config file.')

  logging.basicConfig(level='INFO', format='%(levelname)s %(asctime)s %(name)s] %(message)s')
  app.run(main)
